package com.datastructure.test.distanceknodes;

import java.util.ArrayList;

public class DistanceKNodes {

    public static void main(String[] args) {

    }

    ArrayList<Integer> distanceNodes  = new ArrayList<>();
    int target;
    int distance;
    public ArrayList<Integer> distanceKnodes (TreeNode root, int target, int k) {
        // write code here
        this.target = target;
        this.distance = k;
        findTarget(root);
        return distanceNodes;
    }

    public int findTarget(TreeNode root){
        if(root == null) return -1;
        if(root.val == target){
            getInt(root.left, 1);
            getInt(root.right, 1);
            return 1;
        }
        int i = findTarget(root.left);
        if(i > 0){
            if(i == distance){
                distanceNodes.add(root.val);
            }else if(i < distance){
                getInt(root.right, i + 1);
            }
            return i + 1;
        }
        i = findTarget(root.right);
        if (i == distance) distanceNodes.add(root.val);
        else if (i < distance) getInt(root.left, i + 1);
        return i + 1;
    }

    public void getInt(TreeNode root, int i){
        if(root == null) return;
        if(i == distance){
            distanceNodes.add(root.val);
            return;
        }
        getInt(root.left, i + 1);
        getInt(root.right, i + 1);
    }
}

class TreeNode {
   int val = 0;
   TreeNode left = null;
   TreeNode right = null;
   public TreeNode(int val) {
     this.val = val;
   }
 }
